#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <random>  // 用于生成随机数
#include <iostream>
#include <string>
#include <sys/mman.h>
#include <fcntl.h>
#include <unistd.h>


namespace py = pybind11;


class RunStat
{
public:
    RunStat() {}
    ~RunStat() {}
private:
    static RunStat* p;
public:
    static RunStat* instance()
    {
        return p;
    }
    py::dict shm_pubs;
    py::dict shm_subs;
    std::vector<char*> shm_ptrs;
};


RunStat* RunStat::p = new RunStat;
const int g_shm_size = 1024 * 1024 * 120;


void init_submem(const std::string mem_name)
{
    int shm_fd = shm_open(mem_name.c_str(), O_RDONLY, 0666);
    if (shm_fd == -1)
    {
        std::cerr << "Failed to create shared memory object" << std::endl;
        std::abort();
    }
    char* shm_ptr = (char*)mmap(0, g_shm_size, PROT_READ, MAP_SHARED, shm_fd, 0);
    if (shm_ptr == MAP_FAILED)
    {
        std::cerr << "Failed to map shared memory" << std::endl;
        std::abort();
    }

    int ptr_id = static_cast<int>(RunStat::instance()->shm_ptrs.size());
    RunStat::instance()->shm_ptrs.push_back(shm_ptr);
    RunStat::instance()->shm_subs[mem_name.c_str()] = py::dict();
    RunStat::instance()->shm_subs[mem_name.c_str()]["fd"] = shm_fd;
    RunStat::instance()->shm_subs[mem_name.c_str()]["ptr_id"] = ptr_id;
    std::cout << "init_submem" << std::endl;
}


py::array_t<uint8_t> sms2cvimg_uint8(py::dict dict)
{
    std::string type = dict["type"].cast<std::string>();
    std::string encoding = dict["encoding"].cast<std::string>();
    std::string shm_name = dict["shm_name"].cast<std::string>();
    int h = dict["height"].cast<int>();
    int w = dict["width"].cast<int>();
    int offset = dict["offset"].cast<int>();
    int shm_fst = dict["shm_fst"].cast<int>();

    if (!RunStat::instance()->shm_subs.contains(shm_name.c_str()) || 1 == shm_fst)
    {
        init_submem(shm_name);
    }

    int ptr_id = RunStat::instance()->shm_subs[shm_name.c_str()]["ptr_id"].cast<int>();
    char* shm_ptr = RunStat::instance()->shm_ptrs[ptr_id];
    shm_ptr = shm_ptr + offset;

    int tp = 0;
    if ("8UC1" == encoding)
        tp = 1;
    else if ("8UC2" == encoding)
        tp = 2;
    else if ("8UC3" == encoding)
        tp = 3;
    else if ("8UC4" == encoding)
        tp = 4;
    else
    {
        throw std::runtime_error("Unsupported sms::encoding type!");
    }

    py::array_t<uint8_t> result({h, w, tp});
    py::buffer_info buf = result.request();
    std::memcpy(buf.ptr, shm_ptr, h * w * tp);
    
    if (RunStat::instance()->shm_ptrs[ptr_id] == MAP_FAILED)
    {
        throw std::runtime_error("Failed to map shared memory!");
    }

    return result;
}


py::array_t<uint16_t> sms2cvimg_uint16(py::dict dict)
{
    std::string type = dict["type"].cast<std::string>();
    std::string encoding = dict["encoding"].cast<std::string>();
    std::string shm_name = dict["shm_name"].cast<std::string>();
    int h = dict["height"].cast<int>();
    int w = dict["width"].cast<int>();
    int offset = dict["offset"].cast<int>();
    int shm_fst = dict["shm_fst"].cast<int>();

    if (!RunStat::instance()->shm_subs.contains(shm_name.c_str()) || 1 == shm_fst)
    {
        init_submem(shm_name);
    }

    int ptr_id = RunStat::instance()->shm_subs[shm_name.c_str()]["ptr_id"].cast<int>();
    char* shm_ptr = RunStat::instance()->shm_ptrs[ptr_id];
    shm_ptr = shm_ptr + offset;

    int tp = 0;
    if ("16UC1" == encoding)
        tp = 1;
    else if ("16UC2" == encoding)
        tp = 2;
    else if ("16UC3" == encoding)
        tp = 3;
    else if ("16UC4" == encoding)
        tp = 4;
    else
    {
        throw std::runtime_error("Unsupported sms::encoding type!");
    }

    py::array_t<uint16_t> result({h, w, tp});
    py::buffer_info buf = result.request();
    std::memcpy(buf.ptr, shm_ptr, h * w * tp * 2);
    
    if (RunStat::instance()->shm_ptrs[ptr_id] == MAP_FAILED)
    {
        throw std::runtime_error("Failed to map shared memory!");
    }

    return result;
}


py::array_t<float> sms2cvimg_float(py::dict dict)
{
    std::string type = dict["type"].cast<std::string>();
    std::string encoding = dict["encoding"].cast<std::string>();
    std::string shm_name = dict["shm_name"].cast<std::string>();
    int h = dict["height"].cast<int>();
    int w = dict["width"].cast<int>();
    int offset = dict["offset"].cast<int>();
    int shm_fst = dict["shm_fst"].cast<int>();

    if (!RunStat::instance()->shm_subs.contains(shm_name.c_str()) || 1 == shm_fst)
    {
        init_submem(shm_name);
    }

    int ptr_id = RunStat::instance()->shm_subs[shm_name.c_str()]["ptr_id"].cast<int>();
    char* shm_ptr = RunStat::instance()->shm_ptrs[ptr_id];
    shm_ptr = shm_ptr + offset;

    int tp = 0;
    if ("32FC1" == encoding)
        tp = 1;
    else if ("32FC2" == encoding)
        tp = 2;
    else if ("32FC3" == encoding)
        tp = 3;
    else if ("32FC4" == encoding)
        tp = 4;
    else
    {
        throw std::runtime_error("Unsupported sms::encoding type!");
    }

    py::array_t<float> result({h, w, tp});
    py::buffer_info buf = result.request();
    std::memcpy(buf.ptr, shm_ptr, h * w * tp * 4);
    
    if (RunStat::instance()->shm_ptrs[ptr_id] == MAP_FAILED)
    {
        throw std::runtime_error("Failed to map shared memory!");
    }

    return result;
}



// 定义一个函数，接受 NumPy 数组并对其元素进行平方
py::array_t<double> square_array(py::array_t<double> input) {
    // 获取输入数组的缓冲区信息
    py::buffer_info buf = input.request();

    // 检查输入数组的维度
    if (buf.ndim != 1) {
        throw std::runtime_error("Number of dimensions must be 1");
    }

    // 获取指向数据的指针
    double *ptr = static_cast<double *>(buf.ptr);

    // 创建一个与输入数组大小相同的输出数组
    py::array_t<double> result(buf.size);
    py::buffer_info buf_out = result.request();
    double *ptr_out = static_cast<double *>(buf_out.ptr);

    // 对输入数组的每个元素进行平方，并存储到输出数组中
    for (size_t i = 0; i < buf.size; i++) {
        ptr_out[i] = ptr[i] * ptr[i];
    }

    return result;
}


py::array_t<double> random_matrix() {
    // 创建一个 3x3 的 NumPy 数组
    py::array_t<double> result({3, 3});
    py::buffer_info buf = result.request();
    double *ptr = static_cast<double *>(buf.ptr);

    // 填充随机数到数组中
    for (size_t i = 0; i < 3; i++) {
        for (size_t j = 0; j < 3; j++) {
            ptr[i * 3 + j] = 1.0;  // 按行优先顺序填充
        }
    }

    return result;
}


// 定义 Python 模块
PYBIND11_MODULE(csms_shm, m) {
    m.def("square_array", &square_array, "Square each element of a NumPy array");
    m.def("random_matrix", &random_matrix, "Generate a 3x3 random matrix");
    m.def("sms2cvimg_uint8", &sms2cvimg_uint8, "SpireMS image to cv::Mat(CV_8U) with shared memory");
    m.def("sms2cvimg_uint16", &sms2cvimg_uint16, "SpireMS image to cv::Mat(CV_16U) with shared memory");
    m.def("sms2cvimg_float", &sms2cvimg_float, "SpireMS image to cv::Mat(CV_32F) with shared memory");
}

